import string
import time

import paramiko


class UbuntuSSHClient:
    def __init__(self, ip, username, password, port=22):
        self.ip = ip
        self.username = username
        self.password = password
        self.port = port
        self.ssh = paramiko.SSHClient()
        self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.ssh.connect(ip, port, username, password)
        self.chan = self.ssh.invoke_shell()
        self._init_chan()

    def _init_chan(self):
        self.run(' ')

    def run(self, command):
        res = self._exec(command)
        return res

    def _exec(self, command):
        self.chan.send(command + '\n')
        if 'start' in command:
            time.sleep(15)
        else:
            time.sleep(1)
        res = b''
        while True:
            if self.chan.recv_ready():
                tmp = self.chan.recv(102400)
                res = res + tmp
                time.sleep(0.1)
                if not self.chan.recv_ready() and len(res) > 0:
                    break
        res = self._format_res(res)
        return res

    @staticmethod
    def _format_res(res: bytes):
        res = res.replace(b'\x1b[?2004l', b'')
        res = res.replace(b'\x1b[?2004h', b'')
        res = res.replace(b'\x07', b'')
        res = res.replace(b'\x1b[?2004h', b'')
        res = res.decode('utf8')
        return res

    def close(self):
        self.chan.close()
        self.ssh.close()
